package com.mr.job;

import com.mr.split.InputSplit;
import com.mr.task.MapTask;
import com.mr.task.ReduceTask;

import java.util.*;
import java.util.concurrent.*;
import java.io.File;
import java.io.FileNotFoundException;
import java.io.IOException;
import java.nio.file.Files;
import java.nio.file.Path;
import java.nio.file.Paths;

public class JobTracker {
    // 作业输出目录映射表
    private final Map<String, String> jobOutputDirs = new ConcurrentHashMap<>();

    // 任务状态跟踪
    private final Map<Integer, TaskStatus> mapTaskStatus = new ConcurrentHashMap<>();
    private final Map<Integer, TaskStatus> reduceTaskStatus = new ConcurrentHashMap<>();
    private static final int MAX_RETRY = 3;

    public void run(Job<?, ?, ?, ?, ?, ?> job, String jobId) throws Exception {
        // 线程池创建
        ExecutorService mapExecutor = Executors.newFixedThreadPool(job.getMapThreads());
        int reduceTasks = job.getReduceThreads();
        ExecutorService reduceExecutor = Executors.newFixedThreadPool(reduceTasks);
        
        try {
            // 使用传入的作业ID
            String jobOutputDir = "output/" + jobId + "/";
            job.setOutputDir(jobOutputDir);
            jobOutputDirs.put(jobId, jobOutputDir);

            // 输入文件存在性检查
            Path inputPath = Paths.get(job.getInputPath());
            if (!Files.exists(inputPath)) {
                throw new FileNotFoundException("Input path not found: " + job.getInputPath());
            }
            
            // 1. 输入分片
            List<InputSplit> splits = splitInput(job.getInputPath());
            
            // 2. 提交Map任务（带重试机制）
            List<Future<Map<Integer, String>>> mapFutures = new ArrayList<>();
            for (int i = 0; i < splits.size(); i++) {
                final int taskId = i;
                mapFutures.add(executeWithRetry(() -> {
                    MapTask<?, ?, ?, ?> task = new MapTask<>(splits.get(taskId), job, taskId);
                    return mapExecutor.submit(task);
                }, taskId, true));
            }

            // 3. 聚合Map结果
            Map<Integer, List<String>> reduceInputs = new ConcurrentHashMap<>();
            for (Future<Map<Integer, String>> future : mapFutures) {
                Map<Integer, String> partitionFiles = future.get();
                partitionFiles.forEach((partition, file) -> 
                    reduceInputs.computeIfAbsent(partition, k -> new ArrayList<>()).add(file));
            }

            // 4. 提交Reduce任务
            List<Future<Void>> reduceFutures = new ArrayList<>();
            for (int partition : reduceInputs.keySet()) {
                ReduceTask<?, ?, ?, ?> reduceTask = new ReduceTask<>(reduceInputs.get(partition), job, partition);
                reduceFutures.add(reduceExecutor.submit(reduceTask));
            }

            // 5. 等待所有任务完成
            for (Future<Map<Integer, String>> future : mapFutures) {
                future.get(); // 自动抛出执行异常
            }
            for (Future<Void> future : reduceFutures) {
                future.get();
            }
        } catch (ExecutionException e) {
            throw new Exception("Job failed due to task error: " + e.getCause().getMessage(), e);
        } finally {
            // 关闭当前作业的线程池并清理资源
            mapExecutor.shutdown();
            reduceExecutor.shutdown();
            if (!mapExecutor.awaitTermination(60, TimeUnit.SECONDS)) {
                mapExecutor.shutdownNow();
            }
            if (!reduceExecutor.awaitTermination(60, TimeUnit.SECONDS)) {
                reduceExecutor.shutdownNow();
            }
        }
    }

    // 重试执行方法
    private <T> Future<T> executeWithRetry(Callable<Future<T>> task, int taskId, boolean isMap) {
        for (int retry = 0; retry < MAX_RETRY; retry++) {
            try {
                Future<T> future = task.call();
                updateTaskStatus(taskId, TaskStatus.RUNNING, isMap);
                return future;
            } catch (Exception e) {
                if (retry == MAX_RETRY - 1) {
                    updateTaskStatus(taskId, TaskStatus.FAILED, isMap);
                    throw new RuntimeException("Task failed after " + MAX_RETRY + " retries");
                }
                updateTaskStatus(taskId, TaskStatus.RETRYING, isMap);
            }
        }
        return null;
    }

    // 状态更新方法
    private void updateTaskStatus(int taskId, TaskStatus status, boolean isMap) {
        if (isMap) {
            mapTaskStatus.put(taskId, status);
        } else {
            reduceTaskStatus.put(taskId, status);
        }
    }

    // 任务状态枚举
    private enum TaskStatus {
        PENDING, RUNNING, COMPLETED, RETRYING, FAILED
    }

    // 获取作业输出目录方法
    public String getJobOutputDir(String jobId) {
        return jobOutputDirs.getOrDefault(jobId, "output/");
    }

    // 清理方法
    public void cleanupJob(String jobId) {
        String outputDir = jobOutputDirs.get(jobId);
        if (outputDir != null) {
            try {
                Files.walk(Paths.get(outputDir))
                    .sorted(Comparator.reverseOrder())
                    .map(Path::toFile)
                    .forEach(File::delete);
                jobOutputDirs.remove(jobId);
            } catch (IOException e) {
                System.err.println("清理作业输出目录失败: " + e.getMessage());
            }
        }
    }

    private List<InputSplit> splitInput(String inputPath) {
        List<InputSplit> splits = new ArrayList<>();
        Path path = Paths.get(inputPath);
        long fileSize = path.toFile().length();
        long splitSize = 1024 * 1024; // 每个分片大小为1MB
        
        long start = 0;
        while (start < fileSize) {
            long length = Math.min(splitSize, fileSize - start);
            splits.add(new InputSplit(path, start, length));
            start += length;
        }
        return splits;
    }
}